##########################################################################################

library(data.table)
library(optparse)
library(dplyr)
library(RColorBrewer)
library(ggpubr)
library(ggsci)
library(ggtern)

##########################################################################################

option_list <- list(
    make_option(c("--muti_cancer"), type = "character") ,
    make_option(c("--muti_pre"), type = "character") ,
    make_option(c("--ccf_file"), type = "character") ,
    make_option(c("--type"), type = "character") ,
    make_option(c("--trunk_file"), type = "character") ,
    make_option(c("--gene_list"), type = "character") ,
    make_option(c("--sample_info"), type = "character") ,
    make_option(c("--out_path"), type = "character")
)

if(1!=1){
    
    muti_cancer <- "~/20220915_gastric_multiple/dna_combinePublic/maf/All_GGA.cancer.maf"
    muti_pre <- "~/20220915_gastric_multiple/dna_combinePublic/maf/All_GGA.precancer.maf"
    ccf_file <- "~/20220915_gastric_multiple/dna_combinePublic/mutationTime/result/All_CCF_mutTime.tsv"
    sample_info <- "~/20220915_gastric_multiple/dna_combinePublic/config/tumor_normal.class.list"
    gene_list <- "~/20220915_gastric_multiple/dna_combinePublic/images/selectGCClone/GCClone_gene.all_record.list"
    out_path <- "~/20220915_gastric_multiple/dna_combinePublic/finalPlot/Evolution_Mode"
    type <- "IGC"
    trunk_file <- "~/20220915_gastric_multiple/dna_combinePublic/finalPlot/Evolution_Mode/ITH.combineNormal.IGC_DGC.TrunkDriver.tsv"

}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

gene_list <- opt$gene_list
muti_cancer <- opt$muti_cancer
muti_pre <- opt$muti_pre
sample_info <- opt$sample_info
out_path <- opt$out_path
ccf_file <- opt$ccf_file
type <- opt$type
trunk_file <- opt$trunk_file

###########################################################################################

dir.create(out_path , recursive = T)
col <- c( "#006699","#DDA520"  )

###########################################################################################

dat_mutiCancer <- fread(muti_cancer  ,sep = "\t",  quote="" ,header = T)
dat_mutiPre <- fread(muti_pre  ,sep = "\t",  quote="" ,header = T)
dat_mutiNodule <- rbind(dat_mutiCancer , dat_mutiPre) 

dat_info <- data.frame(fread(sample_info))
dat_gene <- data.frame(fread(gene_list))
dat_ccf <- data.frame(fread(ccf_file))
dat_ith <- data.frame(fread(trunk_file))

###########################################################################################

Variant_Type <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")

###########################################################################################
dat_mutiNodule <- subset(dat_mutiNodule , t_alt_count > 0)

dat_mutiNodule <- data.frame(Hugo_Symbol = dat_mutiNodule$Hugo_Symbol,
    Chromosome = dat_mutiNodule$Chromosome , Start_Position =  dat_mutiNodule$Start_Position , End_Position = dat_mutiNodule$End_Position ,
    Reference_Allele = dat_mutiNodule$Reference_Allele , Tumor_Seq_Allele2 = dat_mutiNodule$Tumor_Seq_Allele2 , 
    Variant_Classification = dat_mutiNodule$Variant_Classification,
    t_ref_count = dat_mutiNodule$t_ref_count , t_alt_count = dat_mutiNodule$t_alt_count ,
    Tumor_Sample_Barcode = dat_mutiNodule$Tumor_Sample_Barcode)

dat_mutiNodule$Location <- paste( dat_mutiNodule$Hugo_Symbol , dat_mutiNodule$Chromosome , dat_mutiNodule$Start_Position , 
    dat_mutiNodule$Reference_Allele , dat_mutiNodule$Tumor_Seq_Allele2 , sep=":"  )

dat_mutiNodule <- subset( dat_mutiNodule , Variant_Classification %in% Variant_Type)

###########################################################################################
## 注释突变的CCF
dat_mutiNodule$mergeCCF <- paste( dat_mutiNodule$Tumor_Sample_Barcode , dat_mutiNodule$Location , sep = ":" )

dat_ccf$mergeCCF <- paste( dat_ccf$Sample , dat_ccf$Hugo_Symbol , dat_ccf$Chr , dat_ccf$Start_Position , 
    dat_ccf$REF , dat_ccf$ALT , sep=":"  )

dat_mutiNodule <- merge( dat_mutiNodule , dat_ccf[,c("mergeCCF" , "CCF_adj")] , by = "mergeCCF" )
dat_mutiNodule <- merge( dat_mutiNodule , dat_info[,c("Tumor" , "Class")] , by.x = "Tumor_Sample_Barcode" , by.y = "Tumor" )

###########################################################################################
## 关注的克隆簇
if(type == "IGC"){
    #dat_info <- subset( dat_info , Type == "IM + IGC" & Type != "IM + IGC + DGC" )
    dat_info <- subset( dat_info , Type == "IM + IGC" | (Type == "IM + IGC + DGC" & Class != "DGC") )
}else if(type == "DGC"){
    dat_info <- subset( dat_info , Type == "IM + DGC" | (Type == "IM + IGC + DGC" & Class != "IGC") )
}else if(type == "IGC_DGC_IGC"){
    dat_info <- subset( dat_info , Type == "IM + IGC + DGC" & Class != "DGC" )
}else if(type == "IGC_DGC_DGC"){
    dat_info <- subset( dat_info , Type == "IM + IGC + DGC" & Class != "IGC" )
}

dat_info_trunkdriver <- subset( dat_info , ID %in% subset( dat_ith , DriverClass_combine=="TrunkDriver" )$ID )
dat_info_notrunkdriver <- subset( dat_info , ID %in% subset( dat_ith , DriverClass_combine=="NoTrunkDriver" )$ID )

###########################################################################################
## 判断基因多少比例出现在Trunk、Pre_Private、Inv_Private
gene_list <- dat_gene$Gene_Symbol
if(type=="IGC"){
    gene_list <- c("TP53" , "APC" , "PIK3CA" , "MICAL2" , "DNAH3")
}else if(type=="DGC"){
    gene_list <- c("TP53" , "APC" , "PIK3CA" , "MICAL2" , "CDH1" , "RHOA" , "SMAD4" , "NRG1")
}else if(type=="IGC_DGC_IGC"){
    gene_list <- c("TP53" , "MICAL2" , "DNAH3")
}else if(type=="IGC_DGC_DGC"){
    gene_list <- c("TP53" , "MICAL2" , "NRG1")
}


clone_t <- 0
computeFunction <- function(dat_info = dat_info , Type = Type){
    result <- c()
    for( gene in gene_list ){

        trunk_driver_num <- 0
        trunk_driver_clone_num <- 0
        trunk_driver_subclone_num <- 0
        share_driver_num <- 0
        evolutionChoose_driver_num <- 0
        private_driver_num <- 0
        private_driver_pre_num <- 0
        private_driver_inv_num <- 0
        
        trunk_driver_sample <- c()
        trunk_driver_clone_sample <- c()
        trunk_driver_subclone_sample <- c()
        share_driver_sample <- c()
        evolutionChoose_driver_sample <- c()
        private_driver_sample <- c()
        private_driver_pre_sample <- c()
        private_driver_inv_sample <- c()

        ## 以人为单位
        for( Sample in unique(dat_info$ID) ){

            tumors <- subset( dat_info , ID == Sample )$Tumor

            ## 确定感兴趣基因的突变
            tmp <- subset( dat_mutiNodule , Tumor_Sample_Barcode %in% tumors & Variant_Classification %in% Variant_Type & Hugo_Symbol == gene )

            ## 以位点判断突变的数量
            tmp_driver <- tmp %>% 
            group_by(Location) %>%
            summarize( MutTumor = length(Tumor_Sample_Barcode) )
            tmp_driver <- data.frame(tmp_driver)

            combine <- unique(subset( dat_info , Tumor %in% tmp$Tumor_Sample_Barcode)$Class)

            ## 若存在突变
            if( nrow(tmp_driver) > 0 ){

                ## 判断哪些变异为share
                tmp$Share <- FALSE
                for( pos in unique(tmp$Location) ){
                    share <- length(which(subset( tmp , Location==pos)$Class %in% c("IM"))) > 0  & 
                    length(which(subset( tmp , Location==pos)$Class %in% c("IGC" , "DGC"))) > 0
                    tmp[tmp$Location==pos,"Share"] <- share
                }

                ## 判断哪些变异为Trunk
                for( pos in unique(tmp$Location) ){
                    share <- nrow(subset( tmp , Location==pos)) == length(tumors)
                    if(share){
                        tmp[tmp$Location==pos,"Share"] <- "Trunk"
                    }
                }

                ## 若基因为IM和GC共享，则判断其为ShareDriver,不管其它突变是否为私有
                if( length(which(tmp$Share!='FALSE')) > 0 ){
                    if( length(which(tmp$Share=='Trunk')) > 0 ){
                        class <- "TrunkDriver"
                        trunk_driver_num <- trunk_driver_num + 1
                        trunk_driver_sample <- c(trunk_driver_sample , Sample )

                        ## 提取Share的突变,判断是否受到进化选择
                        tmp_share <- subset( tmp , Share == 'Trunk' )
                        ## 人工检查和pyclone的结果符合
                        ## ACKR3在JZ585B的一个DGC为主克隆，另一个为亚克隆，其pyclone聚类判定其为亚克隆
                        if( gene == "ACKR3" ){
                            evolution_choose <- median(tmp_share[tmp_share$Class %in% c("IGC" , "DGC"),"CCF_adj"]) >= clone_t
                        }else{
                            evolution_choose <- max(tmp_share[tmp_share$Class %in% c("IGC" , "DGC"),"CCF_adj"]) >= clone_t
                        }
                        
                        if(evolution_choose){
                            evolutionChoose_driver_num <- evolutionChoose_driver_num + 1
                            evolutionChoose_driver_sample <- c(evolutionChoose_driver_sample , Sample )

                            ## 若为Trunk，判断其是否为clone
                            trunk_driver_clone_num <- trunk_driver_clone_num + 1
                            trunk_driver_clone_sample <- c(trunk_driver_clone_sample , Sample )
                        }else{
                            trunk_driver_subclone_num <- trunk_driver_subclone_num + 1
                            trunk_driver_subclone_sample <- c(trunk_driver_subclone_sample , Sample )
                        }

                    }else{
                        class <- "ShareDriver"
                        share_driver_num <- share_driver_num + 1
                        share_driver_sample <- c(share_driver_sample , Sample )

                        ## 提取Share的突变,判断是否受到进化选择
                        tmp_share <- subset( tmp , Share == TRUE )
                        if( gene == "ACKR3" ){
                            evolution_choose <- median(tmp_share[tmp_share$Class %in% c("IGC" , "DGC"),"CCF_adj"]) >= clone_t
                        }else{
                            evolution_choose <- max(tmp_share[tmp_share$Class %in% c("IGC" , "DGC"),"CCF_adj"]) >= clone_t
                        }
                        if(evolution_choose){
                            evolutionChoose_driver_num <- evolutionChoose_driver_num + 1
                            evolutionChoose_driver_sample <- c(evolutionChoose_driver_sample , Sample )
                        }
                    }

                }else{
                    ## 若基因为私有，则判断其为PrivateDriver
                    class <- "PrivateDriver"
                    private_driver_num <- private_driver_num + 1
                    private_driver_sample <- c(private_driver_sample , Sample )

                    ## 判断该基因是否位于癌前分支
                    if( length(which(combine %in% c("IM") )) > 0 ){
                        private_driver_pre_num <- private_driver_pre_num + 1
                        private_driver_pre_sample <- c(private_driver_pre_sample , Sample )
                    }
                    ## 判断该基因是否位于癌分支
                    if( length(which(combine %in% c("IGC" , "DGC") )) > 0 ){
                        private_driver_inv_num <- private_driver_inv_num + 1
                        private_driver_inv_sample <- c(private_driver_inv_sample , Sample )
                    }
                    
                }
            }
        }

       
        tmp <- data.frame( gene = gene , 
            trunk_driver_clone_num = trunk_driver_clone_num ,
            trunk_driver_subclone_num = trunk_driver_subclone_num ,
            trunk_driver_num = trunk_driver_num ,
            share_driver_num = share_driver_num , 
            evolutionChoose_driver_num = evolutionChoose_driver_num ,
            private_driver_num = private_driver_num ,  
            private_driver_pre_num = private_driver_pre_num , private_driver_inv_num = private_driver_inv_num ,
            share_driver_sample = paste0( share_driver_sample , collapse = "," ) , 
            trunk_driver_sample = paste0( trunk_driver_sample , collapse = "," ) , 
            evolutionChoose_driver_sample = paste0( evolutionChoose_driver_sample , collapse = "," ) , 
            private_driver_sample = paste0( private_driver_sample , collapse = "," ) ,
            private_driver_pre_sample = paste0( private_driver_pre_sample , collapse = "," ) ,
            private_driver_inv_sample = paste0( private_driver_inv_sample , collapse = "," )
        )

        result <- rbind( result , tmp )

    }
    result$Type <- Type
    return(result)
}


dat_info <- dat_info_trunkdriver
Type <- "TrunkDriver"
result_trunk <- computeFunction(dat_info = dat_info , Type = Type)

dat_info <- dat_info_notrunkdriver
Type <- "NoTrunkDriver"
result_notrunk <- computeFunction(dat_info = dat_info , Type = Type)

###########################################################################################
## 基因分布堆叠图
dat <- rbind(result_trunk , result_notrunk)

## 基因的顺序
gene_order <- dat[
    order( (dat$trunk_driver_clone_num + dat$trunk_driver_subclone_num + dat$share_driver_num + dat$private_driver_inv_num + dat$private_driver_pre_num) , 
        dat$trunk_driver_clone_num , dat$trunk_driver_num , dat$share_driver_num , dat$private_driver_inv_num  , decreasing=T),"gene"]

dat <- dat[,c("gene" , "trunk_driver_clone_num" , "trunk_driver_subclone_num" , "share_driver_num" , "private_driver_inv_num" , "private_driver_pre_num" , "Type")]
colnames(dat) <- c("gene" , "Trunk Clone" , "Trunk SubClone" , "Share" , "GC branch" , "IM branch" , "Type")
dat <- melt(data.table(dat),id=c("gene","Type"))

dat$variable <- as.character(dat$variable)
dat$variable <- ifelse( dat$variable %in% c("Trunk SubClone" , "Trunk Clone") , "Trunk" , dat$variable )
dat$variable <- ifelse( dat$variable %in% c("Trunk") , "Share" , dat$variable )
dat$variable <- ifelse( dat$variable %in% c("GC branch") & dat$Type == "TrunkDriver" , "IM->GC branch" , dat$variable )
dat$variable <- ifelse( dat$variable %in% c("GC branch") & dat$Type == "NoTrunkDriver" , "Normal->GC branch" , dat$variable )

dat$variable <- factor( dat$variable , levels = c("IM branch" , "Normal->GC branch" , "IM->GC branch" , "Share") , order = T )
dat$gene <- factor( dat$gene , levels = unique(gene_order) , order = T )
dat$Type <- factor( dat$Type , levels = c("TrunkDriver" , "NoTrunkDriver") , order = T )

#col_use <- c( "#E64B35FF" , "#F39B7FFF" , "#4DBBD5FF" , "#00A087FF" , "#3C5488FF" )
#col_use <- col_use[5:1]
#names(col_use) <- c("IM branch" , "GC branch" , "Share" , "Trunk SubClone" , "Trunk Clone")

col_use <- c( "#E64B35FF" , "#00A087FF" , "#4DBBD5B2" , "#3C5488FF" )
col_use <- col_use[4:1]
names(col_use) <- c("IM branch" , "Normal->GC branch" , "IM->GC branch" , "Share")
col_use <- col_use[2:4]

col_use <- c(rgb(red=179,green=34,blue=35,alpha=255,max=255) ,
    rgb(red=247,green=184,blue=71,alpha=255,max=255) ,
    rgb(red=2,green=100,blue=190,alpha=255,max=255) ,
    "#4DAF4A"
    )
col_use <- col_use[4:1]
names(col_use) <- c("IM branch" , "Normal->GC branch" , "IM->GC branch" , "Share")
col_use <- col_use[2:4]

## 构成比
if(type == "IGC"){
    width = 2
    height= 3 
}else if(type == "DGC"){
    width = 3.2
    height= 3.5
}else if(type == "IGC_DGC_IGC"){
    width = 2.5
    height= 4 
}else if(type == "IGC_DGC_DGC"){
    width = 2.5
    height= 4 
}

if(type!="DGC"){
    legend_position <- "none"
}else{
    legend_position <- "bottom"
}


## 判断基因类型
p1 <- ggplot(dat , aes(gene, weight= value , fill=variable))+
    geom_bar(position="stack")+labs(,y="Number of Patients")+
    #facet_grid(cols = vars(Type) , scales = "free_x" ) +
    theme(panel.background = element_rect(fill = NA, colour = "black", size = 1),
        legend.position =legend_position,
        panel.grid.major = element_line(colour=NA),
        legend.text = element_text(size = 6,color="black",face='bold'),
        axis.text.y = element_text(size = 8,color="black",face='bold'),
        axis.title.y = element_text(size = 9,color="black",face='bold'),
        axis.title.x=element_blank(),
        axis.text.x = element_text(size = 10,color="black",face='bold' , angle = 90 ,  vjust = .5, hjust = .5),
        strip.text.x = element_text(size = 7 , face = 'bold'),
        axis.ticks.length = unit(0.2, "cm") ,
        axis.line = element_line(size = 0.5)) +
    guides(fill = guide_legend(reverse=TRUE , title = NULL)) +
    scale_fill_manual(values = col_use)

p <- p1
gp <- ggplotGrob(p)
facet.columns <- gp$layout$l[grepl("panel", gp$layout$name)]
x.var <- sapply(ggplot_build(p)$layout$panel_scales_x,
                function(l) length(l$range$range))
gp$widths[facet.columns] <- gp$widths[facet.columns] * x.var

image_name <- paste0(out_path , "/GeneTrunk.evolutionMode.",type,".pdf")
pdf(image_name , width=width ,height=height)
grid::grid.draw(gp)
dev.off()

###########################################################################################
## 三元图
result <- c()
for( geneN in unique(dat$gene) ){
    tmp <- subset( dat , gene == geneN)
    subset( tmp , variable == "Share" )
    share_num <- sum(subset( tmp , variable == "Share" )$value)
    im_gc_num <- sum(subset( tmp , variable == "IM->GC branch" )$value)
    normal_gc_num <- sum(subset( tmp , variable == "Normal->GC branch" )$value)

    tmp_res <- data.frame( gene = geneN , share = share_num , im_gc = im_gc_num , normal_gc = normal_gc_num)
    result <- rbind( result , tmp_res )
}

data <- result
data$count=rowSums(data[,2:4])
data$count2=rowSums(data[,2:4]) **2
data$colour=rgb(
    data$share/data$count,
    data$im_gc/data$count,
    data$normal_gc/data$count,max=1)

theme_use <- function (base_size = 12, base_family = ""){

    col_use <- c(
        rgb(24,118,188,alpha=255,maxColorValue=255),
        rgb(5,171,151,alpha=255,maxColorValue=255),
        rgb(255,60,49,alpha=255,maxColorValue=255)
        )
    names(col_use) <- c("Share" , "IM->GC branch" , "Normal->GC branch")

    col.T = col_use["Share"]
    col.L = col_use["IM->GC branch"]
    col.R = col_use["Normal->GC branch"]
    col.text = "black"
    col.bg.strip = "gray90"
    col.bg = "white"
    theme_custom(
        base_size = base_size, 
        base_family = base_family, 
        tern.plot.background = NULL, 
        tern.panel.background = col.bg, 
        col.T = col.T, col.L = col.L, col.R = col.R, col.grid.minor = col.bg.strip) + 
    theme(
        text = element_text(color = col.text), 
        strip.background = element_rect(color = col.text, 
        fill = col.bg.strip), 
        strip.text = element_text(color = col.text), 
        tern.axis.arrow.show = TRUE ,
        legend.position = "right" , 
        axis.text=element_blank(), 
        axis.ticks=element_blank()
        )
}

p <- ggtern(data=data,aes(x=im_gc,y=share,z=normal_gc))+
      geom_mask()+
      geom_point(aes(size=count,fill=colour),shape=21,color="black")+
      scale_fill_manual(breaks = unique(data$colour),values=unique(data$colour)) +
      geom_text( aes(label = gene , color = colour ) , size = 6 , hjust = 0, nudge_x = 0.05 ) +
      theme_use() 

image_name <- paste0(out_path , "/GeneTrunk.evolutionMode.",type,".ggtern.pdf")
ggsave(  image_name , p , width = 6 , height = 6 )